import os
import cv2
import numpy as np
import open3d as o3d


def txt2intrinsic(intrinsic_path):
    I = []
    dist = None
    with open(intrinsic_path, "r") as f:
        for line in f:
            values = [float(v) for v in line.strip().split(",")]
            if len(values) == 3:
                I.append(values)
            else:
                dist = values
    I = np.float32(I)
    dist = np.float32(dist)
    return I, dist


def visualize_rgbd(rgb_path, depth_path, intrinsic_mat):
    rgb = o3d.io.read_image(rgb_path)
    depth = o3d.io.read_image(depth_path)
    intrinsic = o3d.camera.PinholeCameraIntrinsic(1280, 720, intrinsic_mat[0, 0], intrinsic_mat[1, 1], intrinsic_mat[0, 2], intrinsic_mat[1, 2])

    rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(rgb, depth, convert_rgb_to_intensity=False, depth_trunc=10.0)
    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, intrinsic)
    o3d.io.write_point_cloud("./rgbd_pcd.ply", pcd)


def _uint16_to_3uint8(img):
    img1 = img.copy()
    img = np.zeros((img1.shape[0], img1.shape[1], 3)).astype(np.uint8)
    img[..., 0] = img1 & 255
    img[..., 1] = (img1 >> 8) & 255
    return img


def _3uint8_to_uint16(img):
    img1 = img.copy()
    img = img1[..., 0].astype(np.uint16) + (img1[..., 1].astype(np.uint16) * 256)
    # img = img.reshape((img1.shape[0], img1.shape[1], 1))
    return img


def get_frame_data(rgb_path, depth_path, idx=0):
    cap = cv2.VideoCapture(rgb_path)
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    cap.set(cv2.CAP_PROP_FOURCC, fourcc)
    suc = cap.isOpened()
    frame_cnt = -1
    while True:
        suc, img = cap.read()
        if not suc:
            break
        frame_cnt += 1
        if frame_cnt == idx:
            rgb = img[:, :, ::-1].astype(np.uint8)
            break
    
    cap = cv2.VideoCapture(depth_path)
    fourcc = cv2.VideoWriter_fourcc('I', '4', '2', '0')
    cap.set(cv2.CAP_PROP_FOURCC, fourcc)
    suc = cap.isOpened()
    frame_cnt = -1
    while True:
        suc, img = cap.read()
        if not suc:
            break
        frame_cnt += 1
        if frame_cnt == idx:
            depth = img
            print(img[500:510, 500:510])
            break
    
    return rgb, depth



def visualize_rgbd_from_video(rgb_path, depth_path, intrinsic_mat, idx=0):
    rgb, depth = get_frame_data(rgb_path, depth_path, idx)
    # depth = o3d.io.read_image("/home/liuyun/HHO/multi-camera/data/20230612_test/decode/0.png")
    rgb = o3d.geometry.Image(rgb)
    depth = o3d.geometry.Image(depth)
    intrinsic = o3d.camera.PinholeCameraIntrinsic(1280, 720, intrinsic_mat[0, 0], intrinsic_mat[1, 1], intrinsic_mat[0, 2], intrinsic_mat[1, 2])

    rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(rgb, depth, convert_rgb_to_intensity=False, depth_trunc=10.0)
    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, intrinsic)
    o3d.io.write_point_cloud("./rgbd_pcd.ply", pcd)


if __name__ == "__main__":
    rgb_path = "/home/liuyun/HHO/multi-camera/data/20230612_test/_d455_camera1_color_image_raw.mp4"
    depth_path = "/home/liuyun/HHO/multi-camera/data/20230612_test/_d455_camera1_aligned_depth_to_color_image_raw.avi"
    intrinsic_path = "./intrinsics/d455_1.txt"

    intrinsic_mat, _ = txt2intrinsic(intrinsic_path)
    visualize_rgbd_from_video(rgb_path, depth_path, intrinsic_mat, idx=0)
